from typing import List

from langchain_chroma import Chroma
from langchain_core.documents import Document
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import PromptTemplate

from langgraph.graph import StateGraph,START
from typing_extensions import TypedDict

from qw_model import qw_model
from Embed_model import BGE_Embed

# -------------------------------------检索器
# 加载知识库
embedding = BGE_Embed()
vectorstore = Chroma(
    embedding_function=embedding,
    persist_directory='./chroma_db'
)

# 创建检索器
retriever = vectorstore.as_retriever(search_kwargs={'k': 3})
# # 检索出的文档
# result = retriever.invoke('抖音云剪多少钱')
# print(result)

class State(TypedDict):
    question: str
    context: List[Document]
    answer: str

def retrieve(state: State):
    retrieve_docs = vectorstore.similarity_search(state["question"])
    return {"context": retrieve_docs}

def generate(state: State):
    # 提示词
    prompt = PromptTemplate(
        template="""你是一个聪明又有亲和力的销售客服，名叫银弓，擅长用轻松、亲切的语气解答用户关于公司业务、价格、产品咨询的问题。你的目标是提供专业、准确的回答，尤其是价格问题，严格按照相应文档回答，风格需要口语化、自然流畅、贴近用户生活，避免书面化和过于正式的表达以及尊称。\n
              这是检索到的文档: \n\n {context} \n\n
              这是用户问题: {question} \n
              "1. 当用户问题不清楚时，主动询问用户具体需求。"
              "2. 回答价格问题，严格按照检索的文档给出，不可自行编造。"
              "3. 当没有用户所说的地区，就按它的省份价格算，当省份也没有那就只能做不限地区的"
              "4. 保持沟通语气亲切、自然，避免过于正式的措辞，多用口语化的表达和轻松的语气，比如“老板”、“宝子”、“咱们”，禁止使用'您'。"
              "5. 输出答案应简洁明了，字数在15字左右。" \n
               需满足以上回答要求。""",
        input_variables=["context", "question"],
    )

    # LLM
    llm = qw_model()

    # Chain
    rag_chain = prompt | llm | StrOutputParser()

    docs_context = "\n\n".join(doc.page_content for doc in state["context"])
    response = rag_chain.invoke({"question": state["question"], "context": docs_context})
    return {"answer": response}


graph = StateGraph(State).add_sequence([retrieve,generate])
graph.add_edge(START, "retrieve")

graph = graph.compile()

res1 = graph.invoke({"question": "实探，江苏，百粉，发小红书多少钱"})
print(res1["answer"])
# res2 = graph.invoke({"question": "我的名字叫什么"})
# print(res2["answer"])


